# Copyright (c) 2022 Presto Labs Pte. Ltd.
# Author: ziyan

import os
import fire
import json
import re
import sys
import subprocess
from datetime import datetime,timedelta
from python.coin.support.feed_tool.feed_cache.app.coin2_fast_feed_runner import get_fastfeed_setting, load_symbol_groups_setting, convert_mea_to_toml_setting, load_subscribers_setting
from xunkemgmt_client.tool.slack_noti import send_to_slack

default_meas = ["Spot.Binance.v1", "Futures.Binance.v1", "Spot.Okex.v5", "Futures.Okex.v5-swap", "Spot.Bybit.v3", "Futures.Bybit.v3-linear"]

def sbatch_cmds(*, cmds, ncpu, driver_dir, job_name, job_ts, dry):
  job_filename = f'{driver_dir}/sbatch_{job_name}_{job_ts}.sh'
  sbatch_out_file = os.path.realpath(f'sbatch_out.{job_ts}.txt')
  with open(job_filename, 'w') as job_file:
    wrapcmds = (
        ["#!/bin/bash"]
        + [f"#SBATCH -o {sbatch_out_file} "]
        + [""]
        + [f"parallel -j {ncpu} <<EOF"]
        + [f"{cmd}" for cmd in cmds] 
        + ["EOF"])
    jobstr = "\n".join(wrapcmds)
    job_file.write(jobstr)
  print(job_filename)
  print(sbatch_out_file)
  cmd_out = f"sbatch -n{ncpu} --mem-per-cpu=4gb --hint=nomultithread --job-name={job_name} --wait --priority=TOP {job_filename}"
  print(cmd_out)
  if not dry:
    start = datetime.now().timestamp()
    error_code = os.system(cmd_out)
    print(f"took {datetime.now().timestamp()-start}s")
    if error_code != 0:
      print(f"error {error_code}")
      return False
    return True


def gen_interval_feed(
    trading_start_time,
    duration,
    resolution,
    ncpu,
    dry = False,
    interval_feed_machine = "main",
    worker = 1,
    meas = None,
    ):
  interval_feed_root = "/remote/iosg/data-2/buckets/feed.derived.realtime.interval_h5"
  job_ts = f'{datetime.now().strftime("%Y%m%d_%H%M")}'
  job_name = f"realtime_regen_cache_symbol_list_{job_ts}"
  driver_dir = os.path.join("regen_cache", job_name)
  os.makedirs(driver_dir, exist_ok=True)
  
  cmds = []
  if meas:
    meas = meas.split(",")
  traget_meas = meas or default_meas
  traget_machines = [
      "feed-05.ap-northeast-1.aws",
      # "feed-02.ap-northeast-2.aws",
      # "feed-05.cn-hongkong.aliyun",
      # "feed-02.eu-west-1.aws",
      # "feed-02.us-east-1.aws",
      # "feed-01.ap-northeast-1.aws.huobi",
  ]
  symbol_groups = load_symbol_groups_setting("./data/coin2/feed/symbol_groups.toml")
  subscribers = load_subscribers_setting("./data/coin2/feed/subscribers.toml")
  date = trading_start_time.split('T')[0]
  resolution_sec = int(resolution[2:-1])*60 if resolution[-1]=='M' else int(resolution[2:-1])*3600
  print(f"resolution: {resolution_sec}s")
  for machine in traget_machines:
    for mea in traget_meas:
      output_path = os.path.join(driver_dir, f'{mea}.json')
      archive_raw_config_path = get_fastfeed_setting(machine, mea, arbitration=True, trading_date=date, long_warmup=True)
      with open(archive_raw_config_path, "r") as f:
        driver = json.load(f)
      driver['feed']['common']['archive']['machine'] = machine
      subscriber = subscribers[convert_mea_to_toml_setting(mea)]
      recipes = [subscriber["default_recipe"]] if "default_recipe" in subscriber.keys() else ["realtime"]
      if "non_tbs" in subscriber["recipes"].keys() and "Futures" in mea:
        recipes.append("non_tbs")
      avaliable_symbol = get_avaliable_symbol(mea)
      for group,symbols in symbol_groups[convert_mea_to_toml_setting(mea)].items():
        output_path = os.path.join(driver_dir, f'{mea}_{group}.json')
        norms = {symbol for symbol in symbols if '-USD' in symbol or '-BUSD' in symbol}
        norms = list(norms & avaliable_symbol)
        exchanges = {}
        for recipe in recipes:
          exchanges[f'{mea}_{recipe}'] = {
            "mea": mea,
            "products": {
              "norms": norms
            },
            "recipe": recipe
          }
        driver['feed']['exchanges'] = exchanges
        driver['feed']['common']['archive']['worker'] = worker
        with open(output_path, "w") as wf:
          json.dump(driver, wf, indent=2)
        cmd = f"""srun -N1 -n1 --exclusive ../coin_binary/coin2-bazel-bin-3.10/cc/appcoin2/support/feed/interval_feed_tool/interval_feed_generator \
--resolution={resolution_sec} \
--driver={output_path} \
--trading_date={trading_start_time} \
--duration={duration} \
--group={group} \
--enable_time=True \
--interval_feed_root={interval_feed_root} \
--interval_feed_machine={interval_feed_machine}"""
        cmds.append(cmd)
        
  add_task_record(job_name,trading_start_time)
  return sbatch_cmds(cmds=cmds, ncpu=ncpu, driver_dir=driver_dir, job_name=job_name, job_ts=job_ts, dry=dry)
  
          
def check_trading_time(time_str):
  return re.search(r'\d{8}T\d{4}',time_str) or re.search(r'\d{8}',time_str)


def check_period_time(time_str):
  return re.search(r'PT\d+[M|H]',time_str)


def get_avaliable_symbol(mea):
  path = f"./data/coin2/product_info/{mea}.json"
  if os.path.exists(path):
    with open(path,"r") as fp:
      pi = json.load(fp)
      return {symbol_info["symbol"] for symbol_info in pi["product_infos"]}


def check_missing_period(trading_start_time):
  path_template = "/remote/iosg/data-2/buckets/feed.derived.realtime.interval_h5/coin/main/PT1M/%s/%s"
  if "T" in trading_start_time:
    check_time = datetime.strptime(trading_start_time,"%Y%m%dT%H%M") - timedelta(minutes=30)
  else :
    check_time = datetime.strptime(trading_start_time,"%Y%m%d") - timedelta(minutes=30)
  
  feed_missing_time = []
  while True:
    all_found = True
    check_time_str = check_time.strftime("%Y%m%dT%H%M")
    for mea in default_meas:
      file_path = path_template % (mea+"_a",check_time_str)
      if not os.path.exists(file_path) or len(os.listdir(file_path)) != 5:
        feed_missing_time.append(check_time_str)
        all_found = False
        break
    if all_found:
      return feed_missing_time
    else :
      check_time = check_time - timedelta(minutes=30)

    
def check_task_record():
  record_path = "./task_record.json"
  running_period = []
  failed_period = []
  ignore_period = [] 
  if os.path.exists(record_path):
    with open(record_path,"r") as fp:
      task_record = json.load(fp)
    for n in range(len(task_record)-1,-1,-1):
      cmd = "sacct --name \"%s\" | head -n 3" % task_record[n]["job_name"]
      status = subprocess.run(cmd, shell=True, capture_output=True, text=True)
      status = status.stdout.split('\n')
      if ('COMPLETED' in status[2]):
        task_record.pop(n)
      elif ('FAILED' in status[2] or len(status)==3):
        failed_period.append((task_record[n]["job_name"],task_record[n]["trading_start_time"]))
        task_record.pop(n)
      else:
        running_period.append(task_record[n]["trading_start_time"])
    print(task_record)
    with open(record_path,"w") as fp:
      json.dump(task_record,fp)
    if len(failed_period) > 3:
      failed_period.sort()
      ignore_period = [t[1] for t in failed_period[:3]]
      slack_msg = f"real time interval feed FAILED to generate:"
      for fail_job_info in failed_period:
        slack_msg = slack_msg + f"\n    job name: {fail_job_info[0]},  trading time: {fail_job_info[1]}"
      slack_msg = slack_msg + f"\n    {ignore_period} will skip generate"   
      send_to_slack(slack_msg, '#coin_feed_noti', 'msg', mention_list=['ziyan'])
  return running_period,ignore_period


def add_task_record(job_name,trading_start_time):
  record_path = "./task_record.json"
  task_record=[]
  if os.path.exists(record_path):
    with open(record_path,"r") as fp:
      task_record = json.load(fp)
  current_task={"job_name":job_name,"trading_start_time":trading_start_time} 
  task_record.append(current_task)
  with open(record_path,"w") as fp:
    json.dump(task_record,fp)
    
    
def main(trading_start_time,
    duration,
    resolution,
    ncpu,
    dry = False,
    interval_feed_machine = "main",
    worker = 1,
    meas = None,):
  
  trading_start_time = str(trading_start_time)
  if not check_trading_time(trading_start_time) :
    print("Wrong trading_start_time format. format:yyyymmddTHHMM or yyyymmdd")
    return
  if not check_period_time(duration):
    print ("Wrong duration format. format:PTnM or PTnH")
    return
  if not check_period_time(resolution):
    print ("Wrong duration format. format:PTnM or PTnH")
    return
  
  trading_start_times = None  
  if meas is None :
    missing_period = check_missing_period(trading_start_time)
    running_period,ignore_period = check_task_record()
    print(f"missing_period: {missing_period}, running_period: {running_period}, failed_task: {ignore_period}")
    trading_start_times = set(missing_period) - set(running_period) - set(ignore_period)
    trading_start_times = list(trading_start_times)
    trading_start_times.sort()
    trading_start_times.append(trading_start_time)
  trading_start_times = trading_start_times or [trading_start_time]
  print(f"generate interval feed trading start at: {trading_start_times}")
  have_error  = False
  for time in trading_start_times:
    res = gen_interval_feed(time,duration,resolution,ncpu,dry,interval_feed_machine,worker,meas)
    have_error = have_error if have_error else not res
  if have_error:
    sys.exit(1)
  
if __name__ == "__main__":
  fire.Fire(main)
